﻿using System.Collections.Generic;
using System.Linq;
using Ast.Core.Pe.Definitions;

namespace Ast.Core.Pe
{
    public class PeModule
    {
        public DosHeader DosHeader { get; set; }

        public NewExeFileHeader NewExeFileHeader { get; set; }

        public IOptionalHeader OptionalHeader { get; set; }

        public PeModuleType ModuleType
        {
            get
            {
                if (OptionalHeader is OptionalHeader32) return PeModuleType.Pe32;
                if (OptionalHeader is OptionalHeader64) return PeModuleType.Pe32Plus;
                return PeModuleType.Unknown;
            }
        }

        public SectionHeader[] SectionHeaders { get; set; }

        public IDictionary<SectionHeader, byte[]> SectionContents { get; set; }

        public bool HasDirectoryEntry(DataDirectoryNumber number)
        {
            return GetDirectoryEntry(number) != null;
        }

        public DataDirectoryEntry? GetDirectoryEntry(DataDirectoryNumber number)
        {
            if (OptionalHeader == null) return null;

            var index = (int) number;
            var entries = OptionalHeader.DataDirectory;
            if (index >= entries.Length) return null;

            var entry = entries[index];
            if (entry.VirtualAddress == 0) return null;

            return entry;
        }

        public DataDirectoryEntry? RelocationDirectoryEntry
        {
            get { return GetDirectoryEntry(DataDirectoryNumber.BaseRelocationTable); }
        }
        public DataDirectoryEntry? ExportDirectoryEntry
        {
            get { return GetDirectoryEntry(DataDirectoryNumber.ExportDirectory); }
        }

        public DataDirectoryEntry? ImportDirectoryEntry
        {
            get { return GetDirectoryEntry(DataDirectoryNumber.ImportDirectory); }
        }

        public DataDirectoryEntry? DelayLoadImportDirectoryEntry
        {
            get { return GetDirectoryEntry(DataDirectoryNumber.DelayLoadImportDescriptors); }
        }

        public SectionHeader? FindSectionContainingRva(uint rva)
        {
            if (SectionHeaders == null || SectionHeaders.Length == 0) return null;
            var retVal = SectionHeaders.FirstOrDefault(sectionHeader => sectionHeader.Contains(rva));
            if (retVal.VirtualAddress == 0) return null;
            return retVal;
        }

        public uint? RvaToFileOffset(uint rva)
        {
            var sectionHeader = FindSectionContainingRva(rva);
            if (sectionHeader.HasValue)
            {
                return rva - sectionHeader.Value.VirtualAddress + sectionHeader.Value.PointerToRawData;
            }
            return null;
        }

        public IList<RelocatedBlock> RelocatedBlocks { get; set; }

        public IList<ImportedModule> ImportedModules { get; set; }

        public IList<ExportedFunction> ExportedFunctions { get; set; }
    }
}